#include <iostream>
#include "captcha_recognize.h"


//int main() {
//    // 创建数据集
//    // D:/CPP/captcha_lstm_cpp/dataset/verify_code_xuexin_all/valid
//
//    auto dataset = CodeDataset("D:/CPP/captcha_lstm_cpp/dataset/verify_code_xuexin_all/valid/")
//        .map(torch::data::transforms::Stack<>());
//
//    // 创建数据加载器
//    auto data_loader = torch::data::make_data_loader<CodeDataset>(
//        std::move(dataset),
//        torch::data::DataLoaderOptions()
//        .batch_size(32)
//        .workers(4)
//    );
//
//    // 测试数据加载器
//    for (auto& batch : *data_loader) {
//        auto images = batch.data;   // [N, C, H, W]
//        auto labels = batch.target; // [N, max_seq_len]
//
//        std::cout << "Images shape: " << images.sizes() << std::endl;
//        std::cout << "Labels shape: " << labels.sizes() << std::endl;
//    }
//
//    return 0;
//}


//int main() {
//    // 创建数据集
//    auto dataset = CodeDataset("D:/CPP/captcha_lstm_cpp/dataset/verify_code_xuexin_all/valid")
//            .map(CollateFn{});
//
//    // 创建数据加载器
//    auto data_loader = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(
//            std::move(dataset),
//            torch::data::DataLoaderOptions().batch_size(32).workers(0).drop_last(false)
//    );
//
//    // 测试数据加载器
//    for (auto& batch : *data_loader) {
//        auto images = batch.data;   // [N, C, H, W]
//        auto labels = batch.target; // [N, max_seq_len]
//
//        std::cout << "Images shape: " << images.sizes() << std::endl;
//        std::cout << "Labels shape: " << labels.sizes() << std::endl;
//    }
//
//    return 0;
//}



//int main() {
//    // 示例输入
//    auto input = torch::randn({2, 3, 100, 200}); // [batch_size, channels, height, width]
//
//    // 创建模型
//    int64_t output_size = 69; // 假设输出类别数为 69
//    auto model = Model(output_size);
//
//    // 前向传播
//    auto output = model->forward(input);
//
//    // 打印输出形状
//    std::cout << "Output shape: " << output.sizes() << std::endl;
//
//    return 0;
//}



int main() {
#if 0
    //const std::string train_path = "D:/CPP/captcha_lstm_cpp/dataset/verify_code_xuexin_all/train";
    //const std::string valid_path = "D:/CPP/captcha_lstm_cpp/dataset/verify_code_xuexin_all/valid";

    const std::string train_path = "C:/Users/SongpingWang/Desktop/verify_code_xuexin/verify_code_xuexin_generate_train2";
    const std::string valid_path = "C:/Users/SongpingWang/Desktop/verify_code_xuexin/verify_code_xuexin_generate_valid2";
#else
    //const std::string train_path = "E:/Desktop/verify_code_xuexin_all/train";
    //const std::string valid_path = "E:/Desktop/verify_code_xuexin_all/valid";

    const std::string train_path = "E:/Desktop/verify_code_xuexin/verify_code_xuexin_generate_train2";
    const std::string valid_path = "E:/Desktop/verify_code_xuexin/verify_code_xuexin_generate_valid2";
#endif
    const std::string save_path = "./";
    const int BATCH_SIZE = 128;
    const int CLASS_NUM = 69;
    const int TOTAL_EPOCH = 150;
    captcha_recognize captcha_recognize(CLASS_NUM);
    captcha_recognize.Start(BATCH_SIZE, CLASS_NUM, TOTAL_EPOCH, train_path, valid_path, save_path);

    return 0;
}
